from typing import Dict, List, Optional

# from ..core.datatypes import Language
from enum import Enum
from .BaseOutputFormatter import BaseOutputFormatter
from ..format_utils import format_str_list_to_a_string


class Language(Enum):
    EN = "en"


class TypeAnswerFormatter(BaseOutputFormatter):
    available_languages = [Language.EN.value]
    patterns = {
        Language.EN.value: "{0} {1} {2}",  # value is a/and label
    }
    sep = "<sep>"

    @classmethod
    def from_relations(
        cls,
        relations: Optional[List[Dict[str, List[str]]]],
        language: str = "en",
        label2token: Dict[str, str] = None,
    ) -> Optional[str]:

        if language not in cls.available_languages:
            raise ValueError(
                f"Expected language to be one of {cls.available_languages}"
            )

        if relations is None:
            return None

        unique_relation_types = list(set([r["type"].lower() for r in relations]))

        if label2token is not None:
            unique_relation_types = [label2token[r] for r in unique_relation_types]

        # answers = [r + " type exists in the sentence." for r in unique_relation_types]
        answers = [r + " type exists in the Sentence" for r in unique_relation_types]

        return format_str_list_to_a_string(answers, cls.sep)


if __name__ == "__main__":
    # print(AnswerFormatter.from_relations({'head': [{'name': 'Lionel Messi'}, {'type': 'footballer'}, {'name': 'Barcelona'}], 'tail': [{'name': 'Lionel Messi'}, {'type': 'footballer'}, {'name': 'Barcelona'}]}, 'en'))

    # Output: 'Lionel Messi is a footballer and Barcelona is a footballer.
    example = [
        {
            "sentence": "John Wilkes Booth , who assassinated President Lincoln , was an actor .",
            "relations": [
                {
                    "head": {
                        "name": "John Wilkes Booth",
                        "type": "People",
                        "pos": [0, 17],
                    },
                    "type": "Kill",
                    "tail": {
                        "name": "President Lincoln",
                        "type": "People",
                        "pos": [37, 54],
                    },
                },
                {
                    "head": {
                        "name": "John Wilkes Booth",
                        "type": "People",
                        "pos": [0, 17],
                    },
                    "type": "Kill",
                    "tail": {
                        "name": "President Lincoln",
                        "type": "People",
                        "pos": [37, 54],
                    },
                },
            ],
        }
    ]
    # print(AnswerFormatter.from_relations(example["relations"], "en"))
    print(TypeAnswerFormatter.from_relations(example[0]["relations"], "en"))

    # print(EntityFormatter.from_relations(example["relations"], "en"))
